"""
Simple DRP-4DVar Integration Example

This example demonstrates the basic integration of DRP-4DVar with realistic
atmospheric data assimilation scenarios using simplified interfaces.
This version works independently without requiring all GSI modules.
"""

using LinearAlgebra
using Statistics
using Random
using Printf
using SparseArrays

# Set up path and load the FourDVar module directly
push!(LOAD_PATH, "../src/FourDVar")
include("../src/FourDVar/FourDVar.jl")
using .FourDVar

# Set random seed for reproducible results
Random.seed!(42)

println("="^80)
println("DRP-4DVar SIMPLE INTEGRATION EXAMPLE")
println("Basic Atmospheric Data Assimilation Demonstration")
println("="^80)

# =============================================================================
# 1. Create Simplified Atmospheric Test Case
# =============================================================================

println("\n1. Setting up Simplified Atmospheric Test Case...")

# Grid configuration
nx, ny, nz = 20, 15, 10
time_window = 4

# Create realistic background atmospheric field
n_3d = nx * ny * nz
n_2d = nx * ny
n_total = 4 * n_3d + n_2d  # u, v, t, q (3D) + ps (2D)

println("  Grid dimensions: $(nx) × $(ny) × $(nz)")
println("  State vector size: $(n_total)")
println("  Time window: $(time_window) hours")

# Generate realistic background field
background_field = zeros(n_total)

# U-wind (m/s) - with jet stream pattern
u_field = reshape(background_field[1:n_3d], nx, ny, nz)
for k in 1:nz, j in 1:ny, i in 1:nx
    sigma = (k - 1) / (nz - 1)
    lat_factor = sin(π * (j - 1) / (ny - 1))
    u_field[i, j, k] = 25.0 * lat_factor * exp(-((sigma - 0.7)/0.15)^2) + randn() * 2.0
end
background_field[1:n_3d] = vec(u_field)

# V-wind (m/s) - smaller magnitude
v_field = reshape(background_field[(n_3d+1):(2*n_3d)], nx, ny, nz)
for k in 1:nz, j in 1:ny, i in 1:nx
    v_field[i, j, k] = randn() * 3.0
end
background_field[(n_3d+1):(2*n_3d)] = vec(v_field)

# Temperature (K) - with realistic lapse rate
t_field = reshape(background_field[(2*n_3d+1):(3*n_3d)], nx, ny, nz)
for k in 1:nz, j in 1:ny, i in 1:nx
    sigma = (k - 1) / (nz - 1)
    t_field[i, j, k] = 288.15 - 65.0 * sigma + randn() * 2.0
end
background_field[(2*n_3d+1):(3*n_3d)] = vec(t_field)

# Specific humidity (kg/kg) - exponential decrease
q_field = reshape(background_field[(3*n_3d+1):(4*n_3d)], nx, ny, nz)
for k in 1:nz, j in 1:ny, i in 1:nx
    sigma = (k - 1) / (nz - 1)
    q_field[i, j, k] = 0.012 * exp(-2.5 * sigma) + abs(randn()) * 0.001
end
background_field[(3*n_3d+1):(4*n_3d)] = vec(q_field)

# Surface pressure (Pa)
ps_field = reshape(background_field[(4*n_3d+1):end], nx, ny)
for j in 1:ny, i in 1:nx
    ps_field[i, j] = 101325.0 + randn() * 300.0
end
background_field[(4*n_3d+1):end] = vec(ps_field)

println("  Background field statistics:")
@printf("    U-wind: %.2f ± %.2f m/s\n", mean(background_field[1:n_3d]), std(background_field[1:n_3d]))
@printf("    Temperature: %.2f ± %.2f K\n", mean(background_field[(2*n_3d+1):(3*n_3d)]), std(background_field[(2*n_3d+1):(3*n_3d)]))
@printf("    Surface pressure: %.1f ± %.1f Pa\n", mean(background_field[(4*n_3d+1):end]), std(background_field[(4*n_3d+1):end]))

# =============================================================================
# 2. Create Background Error Covariance
# =============================================================================

println("\n2. Creating Background Error Covariance...")

# Create realistic background error covariance (simplified)
function create_background_error_covariance(n_total, nx, ny, nz)
    println("  Building background error covariance matrix...")
    
    # Use sparse representation for efficiency
    I_indices = Int[]
    J_indices = Int[]
    values = Float64[]
    
    # Correlation parameters
    h_corr = 3.0  # Horizontal correlation in grid points
    v_corr = 2.0  # Vertical correlation in levels
    
    # Variable error variances
    var_u = 6.25    # (2.5 m/s)²
    var_v = 6.25    # (2.5 m/s)²  
    var_t = 9.0     # (3.0 K)²
    var_q = 1.0e-6  # (0.001 kg/kg)²
    var_ps = 10000.0 # (100 Pa)²
    
    variances = [
        fill(var_u, n_3d);
        fill(var_v, n_3d);
        fill(var_t, n_3d);
        fill(var_q, n_3d);
        fill(var_ps, n_2d)
    ]
    
    # Diagonal elements
    for i in 1:n_total
        push!(I_indices, i)
        push!(J_indices, i)
        push!(values, variances[i])
    end
    
    # Add correlations (simplified structure)
    for var in 1:4  # 3D variables only
        var_offset = (var - 1) * n_3d
        
        for k in 1:nz, j in 1:ny, i in 1:nx
            idx = var_offset + (k-1)*nx*ny + (j-1)*nx + i
            
            # Horizontal correlations
            for dj in -1:1, di in -1:1
                if di == 0 && dj == 0
                    continue
                end
                
                i2, j2 = i + di, j + dj
                if 1 <= i2 <= nx && 1 <= j2 <= ny
                    idx2 = var_offset + (k-1)*nx*ny + (j2-1)*nx + i2
                    
                    dist = sqrt(di^2 + dj^2)
                    correlation = exp(-0.5 * (dist / h_corr)^2)
                    covariance = correlation * sqrt(variances[idx] * variances[idx2])
                    
                    if abs(covariance) > 0.01
                        push!(I_indices, idx)
                        push!(J_indices, idx2)
                        push!(values, covariance)
                    end
                end
            end
            
            # Vertical correlations
            for dk in [-1, 1]
                k2 = k + dk
                if 1 <= k2 <= nz
                    idx2 = var_offset + (k2-1)*nx*ny + (j-1)*nx + i
                    
                    correlation = exp(-0.5 * (abs(dk) / v_corr)^2)
                    covariance = correlation * sqrt(variances[idx] * variances[idx2])
                    
                    if abs(covariance) > 0.01
                        push!(I_indices, idx)
                        push!(J_indices, idx2)
                        push!(values, covariance)
                    end
                end
            end
        end
    end
    
    # Create sparse matrix
    B_sparse = sparse(I_indices, J_indices, values, n_total, n_total)
    
    # Convert to dense for DRP-4DVar (in practice, keep sparse)
    B_dense = Matrix(B_sparse)
    
    # Ensure positive definiteness
    eigenvals = eigvals(B_dense)
    min_eigenval = minimum(eigenvals)
    if min_eigenval <= 0
        println("  Adding regularization for positive definiteness...")
        B_dense += (1e-6 - min_eigenval) * I
    end
    
    println("  Background error matrix created: $(size(B_dense))")
    println("  Sparsity: $(100*(1 - nnz(B_sparse)/prod(size(B_dense))))%")
    
    return B_dense
end

background_error_cov = create_background_error_covariance(n_total, nx, ny, nz)

# =============================================================================
# 3. Create Realistic Observations
# =============================================================================

println("\n3. Creating Realistic Observations...")

# Generate observation locations and values
n_obs_per_time = 100

observations = Dict{Int, Vector{Float64}}()
observation_operators = Dict{Int, Any}()

for time_idx in 1:time_window
    println("  Creating observations for time step $(time_idx)...")
    
    # Generate random observation locations
    obs_locations = zeros(n_obs_per_time, 3)
    for obs_idx in 1:n_obs_per_time
        obs_locations[obs_idx, 1] = 1 + rand() * (nx - 1)  # i coordinate
        obs_locations[obs_idx, 2] = 1 + rand() * (ny - 1)  # j coordinate
        obs_locations[obs_idx, 3] = 1 + rand() * (nz - 1)  # k coordinate
    end
    
    # Create observation operator for this time
    observation_operators[time_idx] = function(state_vector)
        obs_values = zeros(n_obs_per_time)
        
        for obs_idx in 1:n_obs_per_time
            i, j, k = round.(Int, obs_locations[obs_idx, :])
            i = max(1, min(nx, i))
            j = max(1, min(ny, j))  
            k = max(1, min(nz, k))
            
            # Extract temperature at this location (simplified interpolation)
            temp_idx = 2*n_3d + (k-1)*nx*ny + (j-1)*nx + i
            if temp_idx <= length(state_vector)
                obs_values[obs_idx] = state_vector[temp_idx] + randn() * 1.0  # 1K obs error
            else
                obs_values[obs_idx] = 280.0 + randn() * 1.0
            end
        end
        
        return obs_values
    end
    
    # Generate synthetic observations (truth + error)
    true_observations = observation_operators[time_idx](background_field)
    observation_errors = randn(n_obs_per_time) * 2.0  # 2K observation error
    observations[time_idx] = true_observations + observation_errors
    
    @printf("    Time %d: %d observations, mean=%.2f±%.2f K\n", 
           time_idx, n_obs_per_time, mean(observations[time_idx]), std(observations[time_idx]))
end

# =============================================================================
# 4. Create Simple Model Operators
# =============================================================================

println("\n4. Creating Simple Model Operators...")

model_operators = Dict{Int, Any}()

for time_idx in 1:time_window
    # Simple identity model operator (no time evolution)
    # In practice, this would be the tangent linear model
    model_operators[time_idx] = function(state_vector)
        return state_vector  # Identity operator for simplicity
    end
end

println("  Created model operators for $(time_window) time steps")

# =============================================================================
# 5. Run DRP-4DVar Analysis
# =============================================================================

println("\n5. Running DRP-4DVar Analysis...")

# Configure DRP-4DVar
drp4dvar_config = DRP4DVar(
    ensemble_size = 25,
    max_outer_loops = 3,
    max_inner_loops = 50,
    convergence_tolerance = 1e-5,
    time_window = time_window,
    optimizer = "lbfgs",
    use_localization = true,
    localization_radius = 150.0,
    ensemble_inflation = 1.01
)

println("  DRP-4DVar Configuration:")
println("    Ensemble size: $(drp4dvar_config.ensemble_size)")
println("    Time window: $(drp4dvar_config.time_window)")
println("    Optimizer: $(drp4dvar_config.optimizer)")
println("    Localization: $(drp4dvar_config.use_localization)")

# Run analysis
println("\n  Executing DRP-4DVar algorithm...")
start_time = time()

analysis_state, statistics = run_drp4dvar(
    drp4dvar_config,
    background_field,
    background_error_cov,
    observations,
    observation_operators,
    model_operators
)

execution_time = time() - start_time

println("  DRP-4DVar analysis completed!")

# =============================================================================
# 6. Analyze Results
# =============================================================================

println("\n6. Analyzing Results...")

# Basic statistics
analysis_increment = analysis_state - background_field
increment_norm = norm(analysis_increment)
rms_increment = sqrt(mean(analysis_increment.^2))

println("  ANALYSIS STATISTICS:")
@printf("    Execution time: %.2f seconds\n", execution_time)
@printf("    Final cost: %.6e\n", get(statistics, "final_cost", NaN))
@printf("    Converged: %s\n", get(statistics, "converged", false))
@printf("    Total iterations: %d\n", get(statistics, "total_iterations", 0))
@printf("    Analysis increment norm: %.6e\n", increment_norm)
@printf("    RMS analysis increment: %.6e\n", rms_increment)

# Innovation statistics (observation minus analysis)
total_innovation_sq = 0.0
total_obs_count = 0

for (time_idx, obs) in observations
    h_analysis = observation_operators[time_idx](analysis_state)
    innovations = obs - h_analysis
    total_innovation_sq += sum(innovations.^2)
    total_obs_count += length(innovations)
end

rms_innovation = sqrt(total_innovation_sq / total_obs_count)
@printf("    RMS innovation (O-A): %.6e\n", rms_innovation)

# Analyze increments by variable
println("\n  INCREMENT BY VARIABLE:")
u_increment = analysis_increment[1:n_3d]
v_increment = analysis_increment[(n_3d+1):(2*n_3d)]
t_increment = analysis_increment[(2*n_3d+1):(3*n_3d)]
q_increment = analysis_increment[(3*n_3d+1):(4*n_3d)]
ps_increment = analysis_increment[(4*n_3d+1):end]

@printf("    U-wind RMS increment: %.6e m/s\n", sqrt(mean(u_increment.^2)))
@printf("    V-wind RMS increment: %.6e m/s\n", sqrt(mean(v_increment.^2)))
@printf("    Temperature RMS increment: %.6e K\n", sqrt(mean(t_increment.^2)))
@printf("    Humidity RMS increment: %.6e kg/kg\n", sqrt(mean(q_increment.^2)))
@printf("    Surface pressure RMS increment: %.6e Pa\n", sqrt(mean(ps_increment.^2)))

# Performance metrics
state_updates_per_second = n_total * get(statistics, "total_iterations", 1) / execution_time
@printf("    Computational efficiency: %.0f state-updates/second\n", state_updates_per_second)

# =============================================================================
# 7. Validation and Summary
# =============================================================================

println("\n7. Validation Summary...")

# Check analysis quality
analysis_quality_checks = [
    ("Convergence achieved", get(statistics, "converged", false)),
    ("Finite analysis state", all(isfinite.(analysis_state))),
    ("Reasonable increment magnitude", increment_norm > 0 && increment_norm < norm(background_field)),
    ("Positive final cost", get(statistics, "final_cost", Inf) > 0),
    ("Reasonable execution time", execution_time < 300.0),  # 5 minutes
    ("Innovation improvement", rms_innovation < 10.0)  # Reasonable range
]

println("  QUALITY CHECKS:")
all_passed = true
for (check_name, passed) in analysis_quality_checks
    status = passed ? "✓" : "✗"
    println("    $(status) $(check_name)")
    all_passed = all_passed && passed
end

# =============================================================================
# 8. Final Summary
# =============================================================================

println("\n" * "="^80)
println("DRP-4DVAR SIMPLE INTEGRATION SUMMARY")
println("="^80)

println("\nTest Configuration:")
println("  Domain: $(nx)×$(ny)×$(nz) regional grid")
println("  State vector size: $(n_total)")
println("  Time window: $(time_window) hours")
println("  Total observations: $(sum(length(obs) for obs in values(observations)))")
println("  Ensemble size: $(drp4dvar_config.ensemble_size)")

println("\nResults:")
@printf("  Execution time: %.2f seconds\n", execution_time)
@printf("  Convergence: %s\n", get(statistics, "converged", false) ? "Yes" : "No")
@printf("  Final cost: %.3e\n", get(statistics, "final_cost", NaN))
@printf("  Analysis increment norm: %.3e\n", increment_norm)
@printf("  Innovation RMS: %.3e\n", rms_innovation)
@printf("  Computational efficiency: %.0f state-updates/second\n", state_updates_per_second)

println("\nIntegration Status:")
if all_passed
    println("  🎉 SUCCESS: DRP-4DVar integration working correctly!")
    println("  ✓ All quality checks passed")
    println("  ✓ Algorithm converged successfully")
    println("  ✓ Realistic analysis increments produced")
    println("  ✓ Performance within expected range")
else
    println("  ⚠ WARNING: Some quality checks failed")
    println("  Please review the analysis for potential issues")
end

println("\nKey Capabilities Demonstrated:")
println("  ✓ Realistic atmospheric background field generation")
println("  ✓ Multi-variable background error covariance modeling")
println("  ✓ Multi-time observation processing")
println("  ✓ Ensemble-based dimensionality reduction")
println("  ✓ L-BFGS optimization in reduced space")
println("  ✓ Analysis increment computation and validation")

println("\nNext Steps:")
println("  • Integrate with real GSI observation operators")
println("  • Add satellite observation types")
println("  • Implement realistic model operators (WRF/GFS)")
println("  • Optimize for large-scale operational domains")
println("  • Add comprehensive diagnostics and monitoring")

println("\n" * "="^80)
println("Simple DRP-4DVar integration test completed successfully!")
println("="^80)